#include "packet_handler_forward.h"
#include "pack_ip.h"
#include "pack_tcp.h"
#include "pack_udp.h"
#include "forward_util.h"
#include "forward_service.h"
#include "session_factory.h"
#include "common.h"

void packet_handler_forward::on_handle_ip(iphdr* ptr_ip){
	LOGI("on_handle_ip");
}

void packet_handler_forward::on_handle_tcp(iphdr* ptr_ip, tcphdr* ptr_tcp){

	pack_ip p_ip(ptr_ip);
	pack_tcp p_tcp(ptr_tcp);

	// 两端的IP端口
	string src_ip = p_ip.get_saddr();
	string dst_ip = p_ip.get_daddr();
	uint16_t src_port = p_tcp.get_src();
	uint16_t dst_port = p_tcp.get_dst();

	
	// deal with outgoing packet
	if (dst_ip.compare(STUB_IP) != 0)
	{
		sockaddr_in si_src = forward_util::get_sockaddr_in(src_ip.c_str(), src_port);
		sockaddr_in si_dst = forward_util::get_sockaddr_in(dst_ip.c_str(), dst_port);
		forward_service* svc = forward_service::get_instance();

		forward_pair returned_pair = svc->safe_get_session_pair(src_port);
		// 如果不存在会话
		if (!returned_pair.is_valid()){
			// 忽略fin和rst报文
			bool is_close = p_tcp.get_fin() || p_tcp.get_rst();
			if (!is_close){
				forward_pair created_pair(si_src, si_dst);
				session_factory* fac = session_factory::get_instance();
				forward_session* sess = fac->create_session(session_factory::TCP, created_pair);
				svc->safe_add_session(sess);
			}
		}
		
		// 真实目标 ->转发服务
		p_ip.set_daddr(p_ip.get_saddr_uint());
		p_tcp.set_dst(TCP_FORWARD_PORT);

		// 源IP和源端口 -> 桩IP和源端口
		p_ip.set_saddr(STUB_IP);
		p_tcp.set_src(src_port);

		LOGI("on_handle_tcp outgoing %s:%u -> %s:%u => %s:%u -> %s:%u",
				src_ip.c_str(),
				src_port,
				dst_ip.c_str(),
				dst_port,
				p_ip.get_daddr().c_str(),
				p_tcp.get_src(),
				p_ip.get_daddr().c_str(),
				p_tcp.get_dst());
	}
	
	// deal with forwarder ack
	else{

		forward_service* svc = forward_service::get_instance();
		const unsigned short key = dst_port;
		forward_pair pair = svc->safe_get_session_pair(key);
		// 会话存在
		if (pair.is_valid()){

			LOGI("sess != null");
		
			// 目的改为：运行的程序
			p_ip.set_daddr(pair.get_src_addr().c_str());
			p_tcp.set_dst(pair.get_src_port());

			// 源改为：真实目标
			p_ip.set_saddr(pair.get_dst_addr().c_str());
			p_tcp.set_src(pair.get_dst_port());

		}

		LOGI("on_handle_tcp incoming %s:%u -> %s:%u => %s:%u -> %s:%u",
				src_ip.c_str(),
				src_port,
				dst_ip.c_str(),
				dst_port,
				p_ip.get_saddr().c_str(),
				p_tcp.get_src(),
				p_ip.get_daddr().c_str(),
				p_tcp.get_dst());
	}
}

void packet_handler_forward::on_handle_udp(iphdr* ptr_ip, udphdr* ptr_udp){
	LOGI("on_handle_udp");
	pack_ip p_ip(ptr_ip);
	pack_udp p_udp(ptr_udp);

	// 两端的IP端口
	string src_ip = p_ip.get_saddr();
	string dst_ip = p_ip.get_daddr();
	uint16_t src_port = p_udp.get_src();
	uint16_t dst_port = p_udp.get_dst();

	// deal with outgoing packet
	if (dst_ip.compare(STUB_IP) != 0){
		sockaddr_in si_src = forward_util::get_sockaddr_in(src_ip.c_str(), src_port);
		sockaddr_in si_dst = forward_util::get_sockaddr_in(dst_ip.c_str(), dst_port);
		forward_service* svc = forward_service::get_instance();

		forward_pair returned_pair = svc->safe_get_session_pair(src_port);
		if (!returned_pair.is_valid()){
			session_factory* fac = session_factory::get_instance();
			forward_pair created_pair(si_src, si_dst);
			forward_session* sess = fac->create_session(session_factory::UDP, created_pair);
			svc->safe_add_session(sess);
		}

		// 真实目标 ->转发服务
		p_ip.set_daddr(p_ip.get_saddr_uint());
		p_udp.set_dst(UDP_FORWARD_PORT);

		// 源IP和源端口 -> 桩IP和源端口
		p_ip.set_saddr(STUB_IP);
		p_udp.set_src(src_port);

		LOGI("on_handle_udp outgoing %s:%u -> %s:%u => %s:%u -> %s:%u",
				src_ip.c_str(),
				src_port,
				dst_ip.c_str(),
				dst_port,
				p_ip.get_saddr().c_str(),
				p_udp.get_src(),
				p_ip.get_daddr().c_str(),
				p_udp.get_dst());
	}

	// deal with forwarder ack
	else{

		forward_service* svc = forward_service::get_instance();
		const unsigned short key = dst_port;
		forward_pair pair = svc->safe_get_session_pair(key);
		// 会话存在
		if (pair.is_valid()){

			LOGI("sess != null");
		
			// 目的改为：运行的程序
			p_ip.set_daddr(pair.get_src_addr().c_str());
			p_udp.set_dst(pair.get_src_port());

			// 源改为：真实目标
			p_ip.set_saddr(pair.get_dst_addr().c_str());
			p_udp.set_src(pair.get_dst_port());

		}

		LOGI("on_handle_udp incoming %s:%u -> %s:%u => %s:%u -> %s:%u",
				src_ip.c_str(),
				src_port,
				dst_ip.c_str(),
				dst_port,
				p_ip.get_saddr().c_str(),
				p_udp.get_src(),
				p_ip.get_daddr().c_str(),
				p_udp.get_dst());
	}
}
